import os
import json
import mysql.connector
import requests
import time
from datetime import datetime
from concurrent.futures import ThreadPoolExecutor, as_completed
import re

# 确保JSON中所有值都被正确引用
def ensure_quoted(obj):
    if isinstance(obj, dict):
        return {k: ensure_quoted(v) for k, v in obj.items()}
    elif isinstance(obj, list):
        return [ensure_quoted(item) for item in obj]
    elif isinstance(obj, (int, float, bool, type(None))):
        return str(obj)  # 非字符串值转为字符串
    else:
        return obj  # 已经是字符串或其他类型，保持不变

# 读取提示词模板
def read_prompt_template():
    try:
        with open('提示词模版-简化.md', 'r', encoding='utf-8') as file:
            content = file.read()
            # 提取提示词模板部分
            return content
    except Exception as e:
        print(f"读取提示词模板时出错: {e}")
        return None

# 连接数据库
def connect_db():
    db_config = {
        'host': os.getenv('MYSQL_HOST', 'localhost'),
        'user': os.getenv('MYSQL_USER', 'root'),
        'password': os.getenv('MYSQL_PASSWORD', 'neo..6555894'),
        'database': os.getenv('MYSQL_DATABASE', 'sales_chat_mvp')
    }
    
    try:
        conn = mysql.connector.connect(**db_config)
        return conn
    except Exception as e:
        print(f"连接数据库时出错: {e}")
        return None

# 获取所有聊天记录，按用户分组
def get_chats_by_user(limit=None, offset=0):
    conn = connect_db()
    if not conn:
        return None
    
    try:
        cursor = conn.cursor(dictionary=True)
        
        # 查询聊天记录，支持分页
        query = """
        SELECT 
            id,
            nickname,
            customer_nickname,
            chat_history
        FROM 
            chat_history
        WHERE 
            chat_history IS NOT NULL AND chat_history != ''
            and ai_judgment is null
        """
        
        # 添加分页
        if limit is not None:
            query += f" LIMIT {limit} OFFSET {offset}"
        
        cursor.execute(query)
        records = cursor.fetchall()
        
        if not records:
            print(f"数据库中没有聊天记录 (limit={limit}, offset={offset})")
            return None
        
        # 按用户分组
        chats_by_user = {}
        for record in records:
            # 对于个人聊天，使用customer_nickname作为用户标识
            # 对于群组聊天，使用nickname作为用户标识
            user_id = record['customer_nickname'] if record['customer_nickname'] else record['nickname']
            
            if not user_id:  # 跳过没有用户标识的记录
                continue
                
            if user_id not in chats_by_user:
                chats_by_user[user_id] = {
                    'chat_ids': [],
                    'combined_history': '',
                    'nickname': record['nickname']
                }
            
            # 添加聊天ID
            chats_by_user[user_id]['chat_ids'].append(record['id'])
            
            # 合并聊天历史
            if record['chat_history']:
                if chats_by_user[user_id]['combined_history']:
                    chats_by_user[user_id]['combined_history'] += "\n\n" + record['chat_history']
                else:
                    chats_by_user[user_id]['combined_history'] = record['chat_history']
        
        return chats_by_user
        
    except Exception as e:
        print(f"查询数据库时出错: {e}")
        return None
    
    finally:
        if conn.is_connected():
            cursor.close()
            conn.close()

# 调用LLM API进行分析并解析多个JSON部分
def analyze_with_llm(prompt):
    api_url = "https://ark.cn-beijing.volces.com/api/v3/chat/completions"
    api_key = "17f2f1c8-393b-411e-a9e5-54467bdd6523"
    
    headers = {
        "Content-Type": "application/json",
        "Authorization": f"Bearer {api_key}"
    }
    
    payload = {
        "model": "doubao-1-5-pro-32k-250115",
        "messages": [
            {
                "role": "system",
                "content": "You are a helpful assistant."
            },
            {
                "role": "user",
                "content": prompt
            }
        ]
    }
    
    max_retries = 3
    retry_delay = 2  # 秒
    
    for attempt in range(max_retries):
        try:
            response = requests.post(api_url, headers=headers, json=payload)
            response.raise_for_status()  # 检查HTTP错误
            
            result = response.json()
            if 'choices' in result and len(result['choices']) > 0:
                content = result['choices'][0]['message']['content']
                
                # 解析多个JSON部分
                json_parts = extract_multiple_json_parts(content)
                if json_parts and len(json_parts) > 0:
                    # 如果成功提取了多个JSON部分，返回它们
                    return json_parts
                
                # 如果没有成功提取多个部分，尝试作为单个JSON解析
                try:
                    # 尝试直接解析整个内容
                    json_result = json.loads(content)
                    # 应用ensure_quoted函数处理所有值
                    json_result = ensure_quoted(json_result)
                    # 返回单个JSON作为第一部分
                    return [json.dumps(json_result, ensure_ascii=False)]
                except json.JSONDecodeError:
                    # 如果失败，尝试提取```json和```之间的内容
                    start_marker = "```json"
                    end_marker = "```"
                    
                    start_index = content.find(start_marker)
                    if start_index != -1:
                        start_index += len(start_marker)
                        end_index = content.find(end_marker, start_index)
                        if end_index != -1:
                            json_str = content[start_index:end_index].strip()
                            try:
                                # 解析提取的JSON并应用ensure_quoted函数
                                json_result = json.loads(json_str)
                                json_result = ensure_quoted(json_result)
                                return [json.dumps(json_result, ensure_ascii=False)]
                            except json.JSONDecodeError:
                                # 如果解析失败，尝试修复常见的未引用值问题
                                json_str = fix_unquoted_values(json_str)
                                return [json_str]
                    
                    # 如果还是失败，返回原始内容
                    return [content]
            
            print(f"API响应格式不正确: {result}")
            return None
            
        except requests.exceptions.RequestException as e:
            if attempt < max_retries - 1:
                print(f"API请求失败，正在重试 ({attempt+1}/{max_retries}): {e}")
                time.sleep(retry_delay)
            else:
                print(f"API请求失败，已达到最大重试次数: {e}")
                return None
        except Exception as e:
            print(f"处理API响应时出错: {e}")
            return None

# 提取多个JSON部分
def extract_multiple_json_parts(content):
    # 查找所有JSON代码块
    json_blocks = []
    
    # 查找 ```json ... ``` 格式的代码块
    json_pattern = r"```json\s*([\s\S]*?)\s*```"
    matches = re.findall(json_pattern, content)
    for match in matches:
        try:
            # 尝试解析JSON并添加到结果中
            json_obj = json.loads(match)
            json_obj = ensure_quoted(json_obj)
            json_blocks.append(json.dumps(json_obj, ensure_ascii=False))
        except json.JSONDecodeError:
            # 如果解析失败，尝试修复并添加
            fixed_json = fix_unquoted_values(match)
            json_blocks.append(fixed_json)
    
    # 如果没有找到JSON代码块，尝试查找可能的JSON对象
    if not json_blocks:
        # 查找 { ... } 格式的JSON对象
        obj_pattern = r"(\{[\s\S]*?\})"
        matches = re.findall(obj_pattern, content)
        for match in matches:
            try:
                # 尝试解析JSON并添加到结果中
                json_obj = json.loads(match)
                json_obj = ensure_quoted(json_obj)
                json_blocks.append(json.dumps(json_obj, ensure_ascii=False))
            except json.JSONDecodeError:
                # 如果解析失败，尝试修复并添加
                try:
                    fixed_json = fix_unquoted_values(match)
                    json_blocks.append(fixed_json)
                except:
                    continue
    
    return json_blocks

# 修复未引用的值
def fix_unquoted_values(json_str):
    # 处理常见的未引用值，如 未体现, 无, 未知 等
    import re
    
    # 直接使用更激进的修复方式
    # 匹配所有冒号后未被引号包裹的内容
    pattern = r'(:\s*)([^",\{\}\[\]\s][^",\{\}\[\]\n]*?)(\s*?)(,|\}|\n)'
    
    # 排除合法的JSON值（数字、true、false、null）
    def replace_if_not_valid_json(match):
        value = match.group(2).strip()
        # 如果是数字，或者是true、false、null，则不替换
        if (value.isdigit() or 
            value == 'true' or value == 'false' or value == 'null' or
            (value.replace('.', '', 1).isdigit() and value.count('.') <= 1)):  # 浮点数
            return match.group(1) + value + match.group(3) + match.group(4)
        else:
            return match.group(1) + '"' + value + '"' + match.group(3) + match.group(4)
    
    json_str = re.sub(pattern, replace_if_not_valid_json, json_str)
    
    try:
        # 尝试解析修复后的JSON
        json_result = json.loads(json_str)
        json_result = ensure_quoted(json_result)
        return json.dumps(json_result, ensure_ascii=False)
    except json.JSONDecodeError:
        # 如果仍然无法解析，返回尽可能修复的字符串
        return json_str

# 更新数据库中的6个新字段
def update_ai_analysis(chat_ids, analysis_results, force_update=False):
    conn = connect_db()
    if not conn:
        return False
    
    try:
        cursor = conn.cursor(dictionary=True)
        
        # 更新每个聊天记录的6个新字段
        for chat_id in chat_ids:
            # 如果不强制更新，先检查是否已有分析结果
            if not force_update:
                check_query = """
                SELECT ai_judgment FROM chat_history WHERE id = %s AND ai_judgment IS NOT NULL
                """
                cursor.execute(check_query, (chat_id,))
                result = cursor.fetchone()
                if result and result['ai_judgment']:
                    print(f"聊天记录 {chat_id} 已有分析结果，跳过更新")
                    continue
            
            # 准备更新的字段
            update_fields = {}
            
            # 映射每个部分到对应的字段
            field_mapping = [
                ('ai_judgment', 0),  # 第一部分 - 判断结果
                ('ai_sales_stage', 1),  # 第二部分 - 销售阶段识别
                ('ai_customer_profile', 2),  # 第三部分 - 客户画像完善
                ('ai_deal_intention', 3),  # 第四部分 - 成交意向评估
                ('ai_communication_quality', 4),  # 第五部分 - 沟通质量打分
                ('ai_tags', 5)  # 第六部分 - 标签
            ]
            
            # 为每个部分更新对应字段
            for field_name, index in field_mapping:
                if index < len(analysis_results):
                    try:
                        # 尝试解析JSON
                        json_obj = json.loads(analysis_results[index])
                        update_fields[field_name] = json.dumps(json_obj, ensure_ascii=False)
                    except:
                        print(f"解析第{index+1}部分JSON失败，跳过更新{field_name}字段")
            
            # 特殊处理：将销售阶段识别中的当前阶段添加到ai_judgment字段中
            try:
                # 如果有销售阶段识别结果
                if 'ai_sales_stage' in update_fields:
                    sales_stage = json.loads(update_fields['ai_sales_stage'])
                    current_stage = None
                    
                    # 提取当前阶段信息
                    if '销售阶段识别' in sales_stage and '当前阶段' in sales_stage['销售阶段识别']:
                        current_stage = sales_stage['销售阶段识别']['当前阶段']
                    
                    # 如果成功提取到当前阶段，将其添加到ai_judgment中
                    if current_stage and 'ai_judgment' in update_fields:
                        judgment = json.loads(update_fields['ai_judgment'])
                        judgment['当前销售阶段'] = current_stage
                        update_fields['ai_judgment'] = json.dumps(judgment, ensure_ascii=False)
            except Exception as e:
                print(f"处理销售阶段识别时出错: {e}")
            
            # 如果没有任何可更新的字段，跳过
            if not update_fields:
                print(f"聊天记录 {chat_id} 没有有效的分析结果，跳过更新")
                continue
            
            # 构建更新SQL语句
            update_query = "UPDATE chat_history SET "
            update_params = []
            
            for field, value in update_fields.items():
                update_query += f"{field} = %s, "
                update_params.append(value)
            
            # 添加时间字段
            current_time = datetime.now().strftime('%Y-%m-%d %H:%M:%S')
            update_query += "analysis_time = %s, ai_analysis_time = %s "
            update_params.extend([current_time, current_time])
            
            # 添加WHERE子句
            update_query += "WHERE id = %s"
            update_params.append(chat_id)
            
            # 执行更新
            cursor.execute(update_query, update_params)
        
        conn.commit()
        return True
        
    except Exception as e:
        print(f"更新数据库时出错: {e}")
        conn.rollback()
        return False
    
    finally:
        if conn.is_connected():
            cursor.close()
            conn.close()

# 处理单个用户的聊天记录
def process_user_chat(user_id, user_data, prompt_template, force_update=False):
    print(f"正在处理用户 '{user_id}' 的聊天记录...")
    
    # 将聊天历史插入到提示词模板中
    prompt = prompt_template.replace("待输入历史记录", user_data['combined_history'])
    
    # 调用LLM API进行分析
    analysis_results = analyze_with_llm(prompt)
    
    if analysis_results:
        # 验证分析结果
        valid_results = []
        for i, result in enumerate(analysis_results):
            try:
                # 尝试解析JSON
                json.loads(result)
                valid_results.append(result)
                print(f"第{i+1}部分JSON有效")
            except json.JSONDecodeError as e:
                print(f"第{i+1}部分JSON解析错误: {e}")
                # 尝试修复
                try:
                    fixed_result = fix_unquoted_values(result)
                    json.loads(fixed_result)  # 测试是否可以解析
                    valid_results.append(fixed_result)
                    print(f"第{i+1}部分JSON已修复")
                except:
                    print(f"第{i+1}部分JSON无法修复，将跳过")
                    # 保存错误数据到文件
                    error_dir = "error_logs"
                    os.makedirs(error_dir, exist_ok=True)
                    timestamp = datetime.now().strftime('%Y-%m-%dT%H-%M-%S-%f')[:-3] + 'Z'
                    error_file = os.path.join(error_dir, f"json_error_{user_data['chat_ids'][0]}_{timestamp}.json")
                    with open(error_file, 'w', encoding='utf-8') as f:
                        f.write(result)
                    print(f"错误数据已保存到: {error_file}")
        
        # 更新数据库
        if valid_results:
            success = update_ai_analysis(user_data['chat_ids'], valid_results, force_update)
            if success:
                print(f"用户 '{user_id}' 的分析结果已更新到数据库")
                return True
            else:
                print(f"更新用户 '{user_id}' 的分析结果失败")
                return False
        else:
            print(f"用户 '{user_id}' 没有有效的分析结果")
            return False
    else:
        print(f"分析用户 '{user_id}' 的聊天记录失败")
        return False

# 测试LLM API调用，使用提示词模板和数据库中的前三个记录
def test_with_real_data(force_update=False):
    print("开始测试LLM API调用，使用真实数据...")
    
    # 读取提示词模板
    prompt_template = read_prompt_template()
    if not prompt_template:
        print("无法读取提示词模板，测试退出")
        return False
    
    # 连接数据库
    conn = connect_db()
    if not conn:
        print("无法连接数据库，测试退出")
        return False
    
    try:
        cursor = conn.cursor(dictionary=True)
        
        # 查询前三条聊天记录
        query = """
        SELECT 
            id,
            nickname,
            customer_nickname,
            chat_history
        FROM 
            chat_history
        LIMIT 3
        """
        
        cursor.execute(query)
        records = cursor.fetchall()
        
        if not records:
            print("数据库中没有聊天记录，测试退出")
            return False
        
        print(f"获取到 {len(records)} 条聊天记录进行测试")
        
        success_count = 0
        for idx, record in enumerate(records):
            print(f"\n测试记录 {idx+1}/{len(records)}")
            
            # 获取聊天历史
            chat_history = record['chat_history'] or ""
            if not chat_history.strip():
                print(f"记录 {idx+1} 没有聊天历史，跳过")
                continue
            
            # 将聊天历史插入到提示词模板中
            prompt = prompt_template.replace("待输入历史记录", chat_history)
            
            # 调用LLM API进行分析
            print(f"调用LLM API分析记录 {idx+1}...")
            analysis_results = analyze_with_llm(prompt)
            
            if analysis_results:
                print(f"获取到记录 {idx+1} 的LLM分析结果，共 {len(analysis_results)} 部分")
                
                # 验证分析结果
                valid_results = []
                for i, result in enumerate(analysis_results):
                    try:
                        # 尝试解析JSON
                        json.loads(result)
                        valid_results.append(result)
                        print(f"第{i+1}部分JSON有效")
                    except json.JSONDecodeError as e:
                        print(f"第{i+1}部分JSON解析错误: {e}")
                        # 尝试修复
                        try:
                            fixed_result = fix_unquoted_values(result)
                            json.loads(fixed_result)  # 测试是否可以解析
                            valid_results.append(fixed_result)
                            print(f"第{i+1}部分JSON已修复")
                        except:
                            print(f"第{i+1}部分JSON无法修复，将跳过")
                
                # 更新数据库
                if valid_results:
                    # 如果不强制更新，先检查是否已有分析结果
                    if not force_update:
                        check_query = """
                        SELECT ai_judgment FROM chat_history WHERE id = %s AND ai_judgment IS NOT NULL
                        """
                        cursor.execute(check_query, (record['id'],))
                        check_result = cursor.fetchone()
                        if check_result and check_result['ai_judgment']:
                            print(f"记录 {idx+1} 已有分析结果，跳过更新")
                            continue
                    
                    # 准备更新的字段
                    update_fields = {}
                    
                    # 映射每个部分到对应的字段
                    field_mapping = [
                        ('ai_judgment', 0),  # 第一部分 - 判断结果
                        ('ai_sales_stage', 1),  # 第二部分 - 销售阶段识别
                        ('ai_customer_profile', 2),  # 第三部分 - 客户画像完善
                        ('ai_deal_intention', 3),  # 第四部分 - 成交意向评估
                        ('ai_communication_quality', 4),  # 第五部分 - 沟通质量打分
                        ('ai_tags', 5)  # 第六部分 - 标签
                    ]
                    
                    # 为每个部分更新对应字段
                    for field_name, index in field_mapping:
                        if index < len(valid_results):
                            try:
                                # 尝试解析JSON
                                json_obj = json.loads(valid_results[index])
                                update_fields[field_name] = json.dumps(json_obj, ensure_ascii=False)
                            except:
                                print(f"解析第{index+1}部分JSON失败，跳过更新{field_name}字段")
                    
                    # 特殊处理：将销售阶段识别中的当前阶段添加到ai_judgment字段中
                    try:
                        # 如果有销售阶段识别结果
                        if 'ai_sales_stage' in update_fields:
                            sales_stage = json.loads(update_fields['ai_sales_stage'])
                            current_stage = None
                            
                            # 提取当前阶段信息
                            if '销售阶段识别' in sales_stage and '当前阶段' in sales_stage['销售阶段识别']:
                                current_stage = sales_stage['销售阶段识别']['当前阶段']
                            
                            # 如果成功提取到当前阶段，将其添加到ai_judgment中
                            if current_stage and 'ai_judgment' in update_fields:
                                judgment = json.loads(update_fields['ai_judgment'])
                                judgment['当前销售阶段'] = current_stage
                                update_fields['ai_judgment'] = json.dumps(judgment, ensure_ascii=False)
                    except Exception as e:
                        print(f"处理销售阶段识别时出错: {e}")
                    
                    # 如果没有任何可更新的字段，跳过
                    if not update_fields:
                        print(f"记录 {idx+1} 没有有效的分析结果，跳过更新")
                        continue
                    
                    # 构建更新SQL语句
                    update_query = "UPDATE chat_history SET "
                    update_params = []
                    
                    for field, value in update_fields.items():
                        update_query += f"{field} = %s, "
                        update_params.append(value)
                    
                    # 添加时间字段
                    current_time = datetime.now().strftime('%Y-%m-%d %H:%M:%S')
                    update_query += "analysis_time = %s, ai_analysis_time = %s "
                    update_params.extend([current_time, current_time])
                    
                    # 添加WHERE子句
                    update_query += "WHERE id = %s"
                    update_params.append(record['id'])
                    
                    # 执行更新
                    cursor.execute(update_query, update_params)
                    conn.commit()
                    
                    print(f"记录 {idx+1} 的分析结果已更新到数据库")
                    success_count += 1
                else:
                    print(f"记录 {idx+1} 没有有效的分析结果，跳过更新")
            else:
                print(f"分析记录 {idx+1} 失败")
        
        print(f"\n测试完成，成功分析 {success_count}/{len(records)} 条记录")
        return success_count > 0
        
    except Exception as e:
        print(f"测试过程中出错: {e}")
        return False
    
    finally:
        if conn.is_connected():
            cursor.close()
            conn.close()

# 从群组聊天的AI分析中提取客户昵称并更新数据库
def update_group_customer_nicknames():
    """
    从群组聊天的AI分析中提取客户昵称并更新数据库
    如果判断到 is_group 是 1，且 ai_analysis 包含"核心业务流程分析"，
    则获取 ai_analysis 中的"客户昵称"，然后更新到 customer_nickname 字段
    """
    print("开始从群组聊天的AI分析中提取客户昵称...")
    
    conn = connect_db()
    if not conn:
        print("无法连接数据库，更新退出")
        return False
    
    try:
        cursor = conn.cursor(dictionary=True)
        
        # 查询符合条件的群组聊天记录
        query = """
        SELECT 
            id,
            nickname,
            customer_nickname,
            ai_analysis
        FROM 
            chat_history
        WHERE 
            is_group = 1
            AND ai_analysis IS NOT NULL 
            AND ai_analysis != ''
            AND ai_analysis LIKE '%核心业务流程分析%'
        """
        
        cursor.execute(query)
        records = cursor.fetchall()
        
        if not records:
            print("没有找到符合条件的群组聊天记录")
            return False
        
        print(f"找到 {len(records)} 条符合条件的群组聊天记录")
        
        update_count = 0
        for record in records:
            chat_id = record['id']
            nickname = record['nickname']
            current_customer_nickname = record['customer_nickname']
            ai_analysis = record['ai_analysis']
            
            # 尝试解析AI分析结果
            try:
                # 尝试修复可能存在的未引用值问题
                try:
                    analysis_json = json.loads(ai_analysis)
                except json.JSONDecodeError:
                    # 如果解析失败，尝试修复
                    fixed_analysis = fix_unquoted_values(ai_analysis)
                    analysis_json = json.loads(fixed_analysis)
                
                # 直接从第一层获取客户昵称
                customer_nickname = None
                if "客户昵称" in analysis_json:
                    customer_nickname = analysis_json["客户昵称"]
                
                # 如果找到客户昵称且与当前值不同，则更新数据库
                if customer_nickname and customer_nickname != current_customer_nickname:
                    update_cursor = conn.cursor()
                    update_query = """
                    UPDATE chat_history
                    SET customer_nickname = %s
                    WHERE id = %s
                    """
                    
                    update_cursor.execute(update_query, (customer_nickname, chat_id))
                    conn.commit()
                    update_cursor.close()
                    
                    print(f"已更新群组 '{nickname}' 的客户昵称: '{current_customer_nickname or '无'}' -> '{customer_nickname}'")
                    update_count += 1
                elif customer_nickname == current_customer_nickname:
                    print(f"群组 '{nickname}' 的客户昵称已是最新: '{customer_nickname}'")
                else:
                    print(f"未在群组 '{nickname}' 的AI分析中找到有效的客户昵称")
                    
            except json.JSONDecodeError:
                print(f"群组 '{nickname}' 的AI分析结果格式无效")
            except Exception as e:
                print(f"处理群组 '{nickname}' 的AI分析时出错: {e}")
        
        print(f"更新完成，共更新 {update_count}/{len(records)} 条记录的客户昵称")
        return update_count > 0
        
    except Exception as e:
        print(f"更新客户昵称时出错: {e}")
        return False
    
    finally:
        if conn.is_connected():
            cursor.close()
            conn.close()

# 修复数据库中所有ai_analysis字段的JSON格式
def fix_all_ai_analysis_json():
    """
    修复数据库中所有ai_analysis字段的JSON格式问题
    1. 查询所有包含ai_analysis的记录
    2. 尝试解析JSON，如果失败则进行修复
    3. 更新数据库中的记录
    4. 对于ai_analysis不为空但ai_analysis_time为空的记录，将ai_analysis_time设置为当前时间
    """
    print("开始修复数据库中所有ai_analysis字段的JSON格式...")
    
    conn = connect_db()
    if not conn:
        print("无法连接数据库，修复退出")
        return False
    
    try:
        cursor = conn.cursor(dictionary=True)
        
        # 查询所有包含ai_analysis的记录
        query = """
        SELECT 
            id,
            ai_analysis,
            ai_analysis_time
        FROM 
            chat_history
        WHERE 
            ai_analysis IS NOT NULL 
            AND ai_analysis != ''
        """
        
        cursor.execute(query)
        records = cursor.fetchall()
        
        if not records:
            print("没有找到包含ai_analysis的记录")
            return False
        
        print(f"找到 {len(records)} 条包含ai_analysis的记录")
        
        update_count = 0
        error_count = 0
        time_update_count = 0
        
        current_time = datetime.now().strftime('%Y-%m-%d %H:%M:%S')
        
        for record in records:
            chat_id = record['id']
            ai_analysis = record['ai_analysis']
            ai_analysis_time = record['ai_analysis_time']
            
            # 检查是否需要修复JSON
            needs_fix = False
            try:
                # 尝试解析JSON
                json.loads(ai_analysis)
            except json.JSONDecodeError:
                needs_fix = True
            
            # 检查是否需要更新时间
            needs_time_update = ai_analysis_time is None
            
            # 如果需要修复JSON
            if needs_fix:
                try:
                    # 尝试修复JSON
                    fixed_json = fix_unquoted_values(ai_analysis)
                    
                    # 再次尝试解析，确保修复成功
                    try:
                        json_obj = json.loads(fixed_json)
                        # 应用ensure_quoted函数处理所有值
                        json_obj = ensure_quoted(json_obj)
                        fixed_json = json.dumps(json_obj, ensure_ascii=False)
                        
                        # 更新数据库 - 同时更新JSON和时间
                        update_query = """
                        UPDATE chat_history
                        SET ai_analysis = %s,
                            ai_analysis_time = %s
                        WHERE id = %s
                        """
                        
                        update_cursor = conn.cursor()
                        update_cursor.execute(update_query, (fixed_json, current_time, chat_id))
                        conn.commit()
                        update_cursor.close()
                        
                        print(f"已修复记录 {chat_id} 的ai_analysis字段并更新时间")
                        update_count += 1
                        time_update_count += 1
                    except json.JSONDecodeError:
                        print(f"记录 {chat_id} 的ai_analysis字段修复失败，无法解析")
                        error_count += 1
                except Exception as e:
                    print(f"处理记录 {chat_id} 时出错: {e}")
                    error_count += 1
            # 如果只需要更新时间
            elif needs_time_update:
                try:
                    update_query = """
                    UPDATE chat_history
                    SET ai_analysis_time = %s
                    WHERE id = %s
                    """
                    
                    update_cursor = conn.cursor()
                    update_cursor.execute(update_query, (current_time, chat_id))
                    conn.commit()
                    update_cursor.close()
                    
                    print(f"已更新记录 {chat_id} 的ai_analysis_time字段")
                    time_update_count += 1
                except Exception as e:
                    print(f"更新记录 {chat_id} 的时间时出错: {e}")
        
        print(f"修复完成，共修复JSON {update_count}/{len(records)} 条记录，失败 {error_count} 条")
        print(f"共更新时间字段 {time_update_count}/{len(records)} 条记录")
        return update_count > 0 or time_update_count > 0
        
    except Exception as e:
        print(f"修复ai_analysis字段时出错: {e}")
        return False
    
    finally:
        if conn.is_connected():
            cursor.close()
            conn.close()

# 修复特定记录的ai_analysis字段
def fix_specific_ai_analysis_json(chat_ids):
    """
    修复指定记录的ai_analysis字段的JSON格式问题
    
    参数:
    chat_ids -- 需要修复的聊天记录ID列表
    """
    print(f"开始修复指定的 {len(chat_ids)} 条记录的ai_analysis字段...")
    
    conn = connect_db()
    if not conn:
        print("无法连接数据库，修复退出")
        return False
    
    try:
        cursor = conn.cursor(dictionary=True)
        
        update_count = 0
        error_count = 0
        time_update_count = 0
        
        current_time = datetime.now().strftime('%Y-%m-%d %H:%M:%S')
        
        for chat_id in chat_ids:
            # 查询记录
            query = """
            SELECT ai_analysis, ai_analysis_time
            FROM chat_history
            WHERE id = %s
            """
            
            cursor.execute(query, (chat_id,))
            record = cursor.fetchone()
            
            if not record or not record['ai_analysis']:
                print(f"记录 {chat_id} 不存在或ai_analysis为空")
                continue
            
            ai_analysis = record['ai_analysis']
            ai_analysis_time = record['ai_analysis_time']
            
            # 检查是否需要修复JSON
            needs_fix = False
            try:
                # 尝试解析JSON
                json.loads(ai_analysis)
            except json.JSONDecodeError:
                needs_fix = True
            
            # 检查是否需要更新时间
            needs_time_update = ai_analysis_time is None
            
            # 如果需要修复JSON
            if needs_fix:
                try:
                    # 尝试修复JSON
                    fixed_json = fix_unquoted_values(ai_analysis)
                    
                    # 再次尝试解析，确保修复成功
                    try:
                        json_obj = json.loads(fixed_json)
                        # 应用ensure_quoted函数处理所有值
                        json_obj = ensure_quoted(json_obj)
                        fixed_json = json.dumps(json_obj, ensure_ascii=False)
                        
                        # 更新数据库 - 同时更新JSON和时间
                        update_query = """
                        UPDATE chat_history
                        SET ai_analysis = %s,
                            ai_analysis_time = %s
                        WHERE id = %s
                        """
                        
                        update_cursor = conn.cursor()
                        update_cursor.execute(update_query, (fixed_json, current_time, chat_id))
                        conn.commit()
                        update_cursor.close()
                        
                        print(f"已修复记录 {chat_id} 的ai_analysis字段并更新时间")
                        update_count += 1
                        time_update_count += 1
                    except json.JSONDecodeError:
                        print(f"记录 {chat_id} 的ai_analysis字段修复失败，无法解析")
                        error_count += 1
                except Exception as e:
                    print(f"处理记录 {chat_id} 时出错: {e}")
                    error_count += 1
            # 如果只需要更新时间
            elif needs_time_update:
                try:
                    update_query = """
                    UPDATE chat_history
                    SET ai_analysis_time = %s
                    WHERE id = %s
                    """
                    
                    update_cursor = conn.cursor()
                    update_cursor.execute(update_query, (current_time, chat_id))
                    conn.commit()
                    update_cursor.close()
                    
                    print(f"已更新记录 {chat_id} 的ai_analysis_time字段")
                    time_update_count += 1
                except Exception as e:
                    print(f"更新记录 {chat_id} 的时间时出错: {e}")
            else:
                print(f"记录 {chat_id} 的ai_analysis字段格式正确，时间字段已存在，无需修复")
        
        print(f"修复完成，共修复JSON {update_count}/{len(chat_ids)} 条记录，失败 {error_count} 条")
        print(f"共更新时间字段 {time_update_count}/{len(chat_ids)} 条记录")
        return update_count > 0 or time_update_count > 0
        
    except Exception as e:
        print(f"修复ai_analysis字段时出错: {e}")
        return False
    
    finally:
        if conn.is_connected():
            cursor.close()
            conn.close()

# 主函数
def main():
    # 根据命令行参数选择操作
    import sys
    import argparse
    
    parser = argparse.ArgumentParser(description='分析聊天记录并更新数据库')
    parser.add_argument('--force', action='store_true', help='强制更新已有分析结果')
    parser.add_argument('--update-nicknames', action='store_true', help='仅更新群组聊天的客户昵称')
    parser.add_argument('--fix-json', action='store_true', help='修复数据库中所有ai_analysis字段的JSON格式')
    parser.add_argument('--fix-ids', type=str, help='修复指定ID的记录，多个ID用逗号分隔')
    parser.add_argument('--batch-size', type=int, default=20, help='每批处理的记录数量，默认为20')
    parser.add_argument('--offset', type=int, default=0, help='起始偏移量，默认为0')
    args = parser.parse_args()
    
    if args.fix_json:
        # 修复所有记录的JSON格式
        fix_all_ai_analysis_json()
    elif args.fix_ids:
        # 修复指定ID的记录
        chat_ids = [int(id.strip()) for id in args.fix_ids.split(',')]
        fix_specific_ai_analysis_json(chat_ids)
    elif args.update_nicknames:
        # 仅更新群组聊天的客户昵称
        update_group_customer_nicknames()
    else:
        # 处理数据
        process_all_data(args.force, args.batch_size, args.offset)

# 处理所有数据
def process_all_data(force_update=False, batch_size=20, offset=0):
    # 读取提示词模板
    prompt_template = read_prompt_template()
    if not prompt_template:
        print("无法读取提示词模板，程序退出")
        return
    
    # 获取按用户分组的聊天记录，支持分页
    chats_by_user = get_chats_by_user(batch_size, offset)
    if not chats_by_user:
        print(f"没有获取到聊天记录 (batch_size={batch_size}, offset={offset})，程序退出")
        return
    
    print(f"找到 {len(chats_by_user)} 个用户的聊天记录 (batch_size={batch_size}, offset={offset})")
    if force_update:
        print("强制更新模式：将覆盖所有已有的分析结果")
    else:
        print("普通更新模式：将跳过已有分析结果的记录")
    
    # 使用线程池并发处理
    with ThreadPoolExecutor(max_workers=8) as executor:
        futures = {}
        for user_id, user_data in chats_by_user.items():
            future = executor.submit(process_user_chat, user_id, user_data, prompt_template, force_update)
            futures[future] = user_id
        
        # 处理结果
        success_count = 0
        error_count = 0
        json_error_count = 0
        
        for future in as_completed(futures):
            user_id = futures[future]
            try:
                result = future.result()
                if result:
                    success_count += 1
            except json.JSONDecodeError as e:
                # 记录JSON解析错误，但继续处理其他用户
                print(f"处理用户 '{user_id}' 时遇到JSON解析错误: {e}")
                print(f"继续处理其他用户...")
                json_error_count += 1
            except Exception as e:
                print(f"处理用户 '{user_id}' 时出错: {e}")
                error_count += 1
    
    print(f"处理完成，成功分析 {success_count}/{len(chats_by_user)} 个用户的聊天记录")
    if json_error_count > 0:
        print(f"遇到 {json_error_count} 个JSON解析错误，详细信息请查看上方输出和error_logs目录")
    if error_count > 0:
        print(f"遇到 {error_count} 个其他错误，详细信息请查看上方输出")
    
    # 打印下一批次的命令提示
    next_offset = offset + batch_size
    print(f"\n要处理下一批记录，请运行:")
    print(f"python analyze_chat_with_llm.py --batch-size {batch_size} --offset {next_offset}" + (" --force" if force_update else ""))

if __name__ == "__main__":
    main() 